# 导入tushare
import tushare as ts
# 初始化pro接口
pro = ts.pro_api('d3bdfac8ffb0ad7ae6c830bed29e28580b51f91ee085a6866c238d84')

# 拉取数据
df = pro.daily(**{
    "ts_code": "",
    "trade_date": "",
    "start_date": "",
    "end_date": "",
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
print(df)
# 按交易日期升序排列
df = df.sort_values(by="trade_date", ascending=True)

# 计算 MA5（5日均线）和 MA10（10日均线）
df["ma5"] = df["close"].rolling(window=5).mean()
df["ma10"] = df["close"].rolling(window=10).mean()

# 计算 RSI 指标（14日周期）
def compute_rsi(data, period=14):
    delta = data["close"].diff() # 计算每日价格变化
    gain = delta.where(delta > 0, 0) # 取涨幅，跌幅记为0
    loss = -delta.where(delta < 0, 0) # 取跌幅，涨幅记为0

    avg_gain = gain.rolling(window=period, min_periods=1).mean()
    avg_loss = loss.rolling(window=period, min_periods=1).mean()

    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))

    return rsi

df["rsi"] = compute_rsi(df)

# 显示部分数据
print(df[["trade_date", "close", "ma5", "ma10", "rsi"]].tail(10))